import os
from datetime import datetime

import matplotlib.pyplot as plt
import numpy as np
import torch
from sklearn.metrics import mean_absolute_error, mean_squared_error

from config import config
from mydataloader import MyDataLoader
from mymodel import ARIMABaseline, HoltWintersBaseline, MyLSTMModel

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")


def compute_tra(y_true, y_pred):
    """计算趋势反转准确率（Trend Reversal Accuracy）"""
    true_trend = np.sign(np.diff(y_true))
    pred_trend = np.sign(np.diff(y_pred))
    min_len = min(len(true_trend), len(pred_trend))
    correct = np.sum(true_trend[:min_len] == pred_trend[:min_len])
    return correct / min_len if min_len > 0 else 0


def plot_forecasts(history_data, lstm, a, h, labels=['LSTM', 'ARIMA', 'HW']):
    plt.figure(figsize=(14, 6))

    # 绘制历史数据
    plt.plot(range(len(history_data)), history_data, label='History', color='black')

    # 起点（右对齐）
    models = {'LSTM': lstm, 'ARIMA': a, 'HW': h}

    for label in labels:
        pred = models[label]
        pred_start = len(history_data) - len(pred)
        x_range = range(pred_start, len(history_data))
        plt.plot(x_range, pred, label=label, linestyle='--')

    plt.title("Forecast vs History")
    plt.xlabel("Time Step")
    plt.ylabel("Value")
    plt.legend()
    plt.grid(True)
    plt.tight_layout()
    plt.savefig(os.path.join(config.TEMP_FOLDER, f"plot_{datetime.now().strftime('%Y%m%d_%H%M%S')}.png"))
    plt.show()


def plot_metrics_table(history_data, lstm, a, h, labels=['LSTM', 'ARIMA', 'HW']):
    models = {'LSTM': lstm, 'ARIMA': a, 'HW': h}
    metrics_table = []

    for label in labels:
        pred = models[label]
        pred_start = len(history_data) - len(pred)
        true_part = history_data[pred_start:]

        mae = mean_absolute_error(true_part, pred)
        mse = mean_squared_error(true_part, pred)
        rmse = np.sqrt(mse)
        tra = compute_tra(true_part, pred)

        metrics_table.append([label, f"{mae:.4f}", f"{mse:.4f}", f"{rmse:.4f}", f"{tra:.4f}"])

    # 绘制表格图
    fig, ax = plt.subplots(figsize=(6, 2))
    ax.axis('off')
    col_labels = ["Model", "MAE", "MSE", "RMSE", "TRA"]
    table = ax.table(cellText=metrics_table,
                     colLabels=col_labels,
                     cellLoc='center',
                     colLoc='center',
                     loc='center')
    table.auto_set_font_size(False)
    table.set_fontsize(10)
    table.scale(1, 1.5)

    plt.tight_layout()

    plt.savefig(os.path.join(config.TEMP_FOLDER, f"table_{datetime.now().strftime('%Y%m%d_%H%M%S')}.png"))
    plt.show()


if __name__ == "__main__":
    os.makedirs(config.TEMP_FOLDER, exist_ok=True)

    d = MyDataLoader()
    city_name = d.get_available_city_list()[0]
    attr = d.get_available_attributes(city_name)[0]

    ts_data, city_id, attr_id, seqs = d.get_predict_data(city_name, attr)

    trend_data = np.array([item[1] for item in ts_data], dtype=np.float64)
    _, _, history_data = d.get_normed_history_data(city_name, attr)
    history_data = [i[1] for i in history_data]

    print(trend_data.shape)
    print("数据长度:", len(trend_data))  # 应 ≥ 2*seasonal_periods
    print("是否存在NaN:", np.isnan(trend_data).any())
    print("是否存在Inf:", np.isinf(trend_data).any())

    checkpoint = torch.load(os.path.join(config.MODELS_DIR, "geostyle_fit.pth"), map_location=device, weights_only=False)
    model = MyLSTMModel(
        num_cities=d.get_cities_length(),
        num_attrs=d.get_attributes_length()
    ).to(device)
    model.load_state_dict(checkpoint['model'])
    model.eval()
    hw = HoltWintersBaseline(
        seasonal_periods=4,  # 假设周周期
        trend='add',  # 强制使用加法模型
        seasonal='add',  # 强制加法季节性
    )
    arima = ARIMABaseline(use_auto_arima=True)

    lstm, a, h = [], [], []

    for i, seq in enumerate(seqs):
        input_tensor = seq.unsqueeze(0).to(device)
        city_tensor = torch.tensor([city_id], dtype=torch.long).to(device)
        attr_tensor = torch.tensor([attr_id], dtype=torch.long).to(device)

        with torch.no_grad():
            delta = model(input_tensor, city_tensor, attr_tensor).cpu().numpy()[0][0]

        last_value = seq[-1][0]
        predicted_value = last_value + delta
        lstm.append(float(predicted_value))

        X = np.array([float(i[0]) for i in seq], dtype=np.float64)
        hw.fit(X)
        arima.fit(X)

        a.append(float(arima.predict()[0]))
        h.append(float(hw.predict().item()))

    lstm = d.denormalize_trend_list(city_name, attr, lstm)
    a = d.denormalize_trend_list(city_name, attr, a)
    h = d.denormalize_trend_list(city_name, attr, h)
    history_data = d.denormalize_trend_list(city_name, attr, history_data)

    print(lstm)
    print(a)
    print(h)
    print(history_data)

    plot_forecasts(history_data, lstm, a, h)
    plot_metrics_table(history_data, lstm, a, h)
